// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <memory>
#include <string>

#include <ktest/ktest.h>

#include <nebula/array/array_base.h>
#include <nebula/array/array_dict.h>
#include <nebula/array/array_nested.h>
#include <nebula/array/data.h>
#include <nebula/core/extension_type.h>

#include <turbo/utility/status.h>
#include <nebula/testing/ktest_util.h>
#include <nebula/types/type.h>
#include <nebula/util/endian.h>
#include <turbo/log/logging.h>

namespace nebula {

    void CheckView(const std::shared_ptr<Array> &input,
                   const std::shared_ptr<DataType> &view_type,
                   const std::shared_ptr<Array> &expected) {
        ASSERT_OK_AND_ASSIGN(auto result, input->View(view_type));
        ASSERT_OK(result->validate_full());
        AssertArraysEqual(*expected, *result);
    }

    void CheckView(const std::shared_ptr<Array> &input,
                   const std::shared_ptr<Array> &expected_view) {
        CheckView(input, expected_view->type(), expected_view);
    }

    void CheckViewFails(const std::shared_ptr<Array> &input,
                        const std::shared_ptr<DataType> &view_type) {
        ASSERT_RAISES(turbo::StatusCode::kInvalidArgument, input->View(view_type));
    }

    class IPv4Type : public ExtensionType {
    public:
        IPv4Type() : ExtensionType(fixed_size_binary(4)) {}

        std::string extension_name() const override { return "ipv4"; }

        bool ExtensionEquals(const ExtensionType &other) const override {
            return other.extension_name() == this->extension_name();
        }

        std::shared_ptr<Array> make_array(std::shared_ptr<ArrayData> data) const override {
            DKCHECK_EQ(data->type->id(), Type::EXTENSION);
            DKCHECK_EQ("ipv4", static_cast<const ExtensionType &>(*data->type).extension_name());
            return std::make_shared<ExtensionArray>(data);
        }

        turbo::Result<std::shared_ptr<DataType>> deserialize(
                std::shared_ptr<DataType> storage_type,
                const std::string &serialized) const override {
            return turbo::unimplemented_error("IPv4Type::deserialize");
        }

        std::string serialize() const override { return ""; }
    };

    TEST(TestArrayView, IdentityPrimitive) {
        auto arr = assert_array_from_json(int16(), "[0, -1, 42]");
        CheckView(arr, arr->type(), arr);
        arr = assert_array_from_json(int16(), "[0, -1, 42, null]");
        CheckView(arr, arr->type(), arr);
        arr = assert_array_from_json(boolean(), "[true, false, null]");
        CheckView(arr, arr->type(), arr);
    }

    TEST(TestArrayView, IdentityNullType) {
        auto arr = assert_array_from_json(null(), "[null, null, null]");
        CheckView(arr, arr->type(), arr);
    }

    TEST(TestArrayView, PrimitiveAsPrimitive) {
        auto arr = assert_array_from_json(int16(), "[0, -1, 42]");
        auto expected = assert_array_from_json(uint16(), "[0, 65535, 42]");
        CheckView(arr, expected);
        CheckView(expected, arr);

        arr = assert_array_from_json(int32(), "[0, 1069547520, -1071644672, null]");
        expected = assert_array_from_json(float32(), "[0.0, 1.5, -2.5, null]");
        CheckView(arr, expected);

        arr = assert_array_from_json(timestamp(TimeUnit::SECOND),
                                     R"(["1970-01-01","2000-02-29","3989-07-14","1900-02-28"])");
        expected = assert_array_from_json(int64(), "[0, 951782400, 63730281600, -2203977600]");
        CheckView(arr, expected);
        CheckView(expected, arr);
    }

    TEST(TestArrayView, PrimitiveAsFixedSizeBinary) {
#if NEBULA_LITTLE_ENDIAN
        auto arr = assert_array_from_json(int32(), "[2020568934, 2054316386, null]");
#else
        auto arr = assert_array_from_json(int32(), "[1718579064, 1650553466, null]");
#endif
        auto expected = assert_array_from_json(fixed_size_binary(4), R"(["foox", "barz", null])");
        CheckView(arr, expected);
        CheckView(expected, arr);
    }

    TEST(TestArrayView, StringAsBinary) {
        auto arr = assert_array_from_json(utf8(), R"(["foox", "barz", null])");
        auto expected = assert_array_from_json(binary(), R"(["foox", "barz", null])");
        CheckView(arr, expected);
        CheckView(expected, arr);
    }

    TEST(TestArrayView, StringViewAsBinaryView) {
        for (auto json: {
                R"(["foox", "barz", null])",
                R"(["foox", "barz_not_inlined", null])",
        }) {
            auto arr = assert_array_from_json(utf8_view(), json);
            auto expected = assert_array_from_json(binary_view(), json);
            CheckView(arr, expected);
            CheckView(expected, arr);
        }
    }

    TEST(TestArrayView, StringViewAsBinaryViewInStruct) {
        auto padl = assert_array_from_json(list(int16()), "[[0, -1], [], [42]]");
        auto padr = assert_array_from_json(utf8(), R"(["foox", "barz", null])");

        for (auto json: {
                R"(["foox", "barz", null])",
                R"(["foox", "barz_not_inlined", null])",
        }) {
            auto arr =
                    StructArray::create({padl, assert_array_from_json(utf8_view(), json), padr}, {"", "", ""})
                            .value_or_die();
            auto expected =
                    StructArray::create({padl, assert_array_from_json(binary_view(), json), padr}, {"", "", ""})
                            .value_or_die();

            CheckView(arr, expected);
            CheckView(expected, arr);
        }
    }

    TEST(TestArrayView, PrimitiveWrongSize) {
        auto arr = assert_array_from_json(int16(), "[0, -1, 42]");
        CheckViewFails(arr, int8());
        CheckViewFails(arr, fixed_size_binary(3));
        CheckViewFails(arr, null());
    }

    TEST(TestArrayView, StructAsStructSimple) {
        auto ty1 = STRUCT({field("a", int8()), field("b", int32())});
        auto ty2 = STRUCT({field("c", uint8()), field("d", float32())});

        auto arr = assert_array_from_json(ty1, "[[0, 0], [1, 1069547520], [-1, -1071644672]]");
        auto expected = assert_array_from_json(ty2, "[[0, 0], [1, 1.5], [255, -2.5]]");
        CheckView(arr, expected);
        CheckView(expected, arr);

        // With nulls
        arr = assert_array_from_json(ty1, "[[0, 0], null, [-1, -1071644672]]");
        expected = assert_array_from_json(ty2, "[[0, 0], null, [255, -2.5]]");
        CheckView(arr, expected);
        CheckView(expected, arr);

        // With nested nulls
        arr = assert_array_from_json(ty1, "[[0, null], null, [-1, -1071644672]]");
        expected = assert_array_from_json(ty2, "[[0, null], null, [255, -2.5]]");
        CheckView(arr, expected);
        CheckView(expected, arr);

        ty2 = STRUCT({field("c", uint8()), field("d", fixed_size_binary(4))});
#if NEBULA_LITTLE_ENDIAN
        arr = assert_array_from_json(ty1, "[[0, null], null, [-1, 2020568934]]");
#else
        arr = assert_array_from_json(ty1, "[[0, null], null, [-1, 1718579064]]");
#endif
        expected = assert_array_from_json(ty2, R"([[0, null], null, [255, "foox"]])");
        CheckView(arr, expected);
        CheckView(expected, arr);
    }

    TEST(TestArrayView, StructAsStructNonNullable) {
        auto ty1 = STRUCT({field("a", int8()), field("b", int32())});
        auto ty2 = STRUCT({field("c", uint8(), /*nullable=*/false), field("d", float32())});

        auto arr = assert_array_from_json(ty1, "[[0, 0], [1, 1069547520], [-1, -1071644672]]");
        auto expected = assert_array_from_json(ty2, "[[0, 0], [1, 1.5], [255, -2.5]]");
        CheckView(arr, expected);
        CheckView(expected, arr);

        // With nested nulls
        arr = assert_array_from_json(ty1, "[[0, null], [-1, -1071644672]]");
        expected = assert_array_from_json(ty2, "[[0, null], [255, -2.5]]");
        CheckView(arr, expected);
        CheckView(expected, arr);

        // Nested null cannot be viewed as non-null field
        arr = assert_array_from_json(ty1, "[[0, null], [null, -1071644672]]");
        CheckViewFails(arr, ty2);
    }

    TEST(TestArrayView, StructAsStructWrongLayout) {
        auto ty1 = STRUCT({field("a", int8()), field("b", int32())});
        auto arr = assert_array_from_json(ty1, "[[0, 0], [1, 1069547520], [-1, -1071644672]]");

        auto ty2 = STRUCT({field("c", int16()), field("d", int32())});
        CheckViewFails(arr, ty2);
        ty2 = STRUCT({field("c", int32()), field("d", int8())});
        CheckViewFails(arr, ty2);
        ty2 = STRUCT({field("c", int8())});
        CheckViewFails(arr, ty2);
        ty2 = STRUCT({field("c", fixed_size_binary(5))});
        CheckViewFails(arr, ty2);
    }

    TEST(TestArrayView, StructAsStructWithNullType) {
        auto ty1 = STRUCT({field("a", int8()), field("b", null())});
        auto ty2 = STRUCT({field("c", uint8()), field("d", null())});

        auto arr = assert_array_from_json(ty1, "[[0, null], [1, null], [-1, null]]");
        auto expected = assert_array_from_json(ty2, "[[0, null], [1, null], [255, null]]");
        CheckView(arr, expected);
        CheckView(expected, arr);

        // With nulls and nested nulls
        arr = assert_array_from_json(ty1, "[null, [null, null], [-1, null]]");
        expected = assert_array_from_json(ty2, "[null, [null, null], [255, null]]");
        CheckView(arr, expected);
        CheckView(expected, arr);

        // Moving the null types around
        ty2 = STRUCT({field("c", null()), field("d", uint8())});
        expected = assert_array_from_json(ty2, "[null, [null, null], [null, 255]]");
        CheckView(arr, expected);
        CheckView(expected, arr);

        // Removing the null type
        ty2 = STRUCT({field("c", uint8())});
        expected = assert_array_from_json(ty2, "[null, [null], [255]]");
        CheckView(arr, expected);
        CheckView(expected, arr);
    }

    TEST(TestArrayView, StructAsFlat) {
        auto ty1 = STRUCT({field("a", int16())});
        auto arr = assert_array_from_json(ty1, "[[0], [1], [-1]]");
        auto expected = assert_array_from_json(uint16(), "[0, 1, 65535]");
        CheckView(arr, expected);
        CheckView(expected, arr);

        // With nulls
        arr = assert_array_from_json(ty1, "[[0], null, [-1]]");
        expected = assert_array_from_json(uint16(), "[0, null, 65535]");
        //   CheckView(arr, expected);  // XXX currently fails
        CheckView(expected, arr);

        // With nested nulls => fails
        arr = assert_array_from_json(ty1, "[[0], [null], [-1]]");
        CheckViewFails(arr, uint16());
    }

    TEST(TestArrayView, StructAsFlatWithNullType) {
        auto ty1 = STRUCT({field("a", null()), field("b", int16()), field("c", null())});
        auto arr = assert_array_from_json(ty1, "[[null, 0, null], [null, -1, null]]");
        auto expected = assert_array_from_json(uint16(), "[0, 65535]");
        CheckView(arr, expected);
        CheckView(expected, arr);

        // With nulls
        arr = assert_array_from_json(ty1, "[[null, 0, null], null, [null, -1, null]]");
        expected = assert_array_from_json(uint16(), "[0, null, 65535]");
        //   CheckView(arr, expected);  // XXX currently fails
        CheckView(expected, arr);

        // With nested nulls => fails
        arr = assert_array_from_json(ty1, "[[null, null, null]]");
        CheckViewFails(arr, uint16());
    }

    TEST(TestArrayView, StructAsStructNested) {
        // Nesting tree shape need not be identical
        auto ty1 = STRUCT({field("a", STRUCT({field("b", int8())})), field("d", int32())});
        auto ty2 = STRUCT({field("a", uint8()), field("b", STRUCT({field("b", float32())}))});
        auto arr = assert_array_from_json(ty1, "[[[0], 1069547520], [[-1], -1071644672]]");
        auto expected = assert_array_from_json(ty2, "[[0, [1.5]], [255, [-2.5]]]");
        CheckView(arr, expected);
        CheckView(expected, arr);

        // With null types
        ty1 = STRUCT({field("a", STRUCT({field("xx", null()), field("b", int8())})),
                       field("d", int32())});
        ty2 = STRUCT({field("a", uint8()),
                       field("b", STRUCT({field("b", float32()), field("xx", null())}))});
        arr = assert_array_from_json(ty1, "[[[null, 0], 1069547520], [[null, -1], -1071644672]]");
        expected = assert_array_from_json(ty2, "[[0, [1.5, null]], [255, [-2.5, null]]]");
        CheckView(arr, expected);
        CheckView(expected, arr);

        // XXX With nulls (currently fails)
    }

    TEST(TestArrayView, ListAsListSimple) {
        auto arr = assert_array_from_json(list(int16()), "[[0, -1], [], [42]]");
        auto expected = assert_array_from_json(list(uint16()), "[[0, 65535], [], [42]]");
        CheckView(arr, expected);
        CheckView(expected, arr);

        // With nulls
        arr = assert_array_from_json(list(int16()), "[[0, -1], null, [42]]");
        expected = assert_array_from_json(list(uint16()), "[[0, 65535], null, [42]]");
        CheckView(arr, expected);
        CheckView(expected, arr);

        // With nested nulls
        arr = assert_array_from_json(list(int16()), "[[0, -1], null, [null, 42]]");
        expected = assert_array_from_json(list(uint16()), "[[0, 65535], null, [null, 42]]");
        CheckView(arr, expected);
        CheckView(expected, arr);
    }

    TEST(TestArrayView, FixedSizeListAsFixedSizeList) {
        auto ty1 = fixed_size_list(int16(), 3);
        auto ty2 = fixed_size_list(uint16(), 3);
        auto arr = assert_array_from_json(ty1, "[[0, -1, 42], [5, 6, -16384]]");
        auto expected = assert_array_from_json(ty2, "[[0, 65535, 42], [5, 6, 49152]]");
        CheckView(arr, expected);
        CheckView(expected, arr);

        // With nested nulls
        arr = assert_array_from_json(ty1, "[[0, -1, null], null, [5, 6, -16384]]");
        expected = assert_array_from_json(ty2, "[[0, 65535, null], null, [5, 6, 49152]]");
        CheckView(arr, expected);
        CheckView(expected, arr);
    }

    TEST(TestArrayView, FixedSizeListAsFlat) {
        auto ty1 = fixed_size_list(int16(), 3);
        auto arr = assert_array_from_json(ty1, "[[0, -1, 42], [5, 6, -16384]]");
        auto expected = assert_array_from_json(uint16(), "[0, 65535, 42, 5, 6, 49152]");
        CheckView(arr, expected);
        // CheckView(expected, arr);  // XXX currently fails

        // XXX With nulls (currently fails)
    }

    TEST(TestArrayView, FixedSizeListAsFixedSizeBinary) {
        auto ty1 = fixed_size_list(int32(), 1);
#if NEBULA_LITTLE_ENDIAN
        auto arr = assert_array_from_json(ty1, "[[2020568934], [2054316386]]");
#else
        auto arr = assert_array_from_json(ty1, "[[1718579064], [1650553466]]");
#endif
        auto expected = assert_array_from_json(fixed_size_binary(4), R"(["foox", "barz"])");
        CheckView(arr, expected);
    }

    TEST(TestArrayView, SparseUnionAsStruct) {
        auto child1 = assert_array_from_json(int16(), "[0, -1, 42]");
        auto child2 = assert_array_from_json(int32(), "[0, 1069547520, -1071644672]");
        auto indices = assert_array_from_json(int8(), "[0, 0, 1]");
        ASSERT_OK_AND_ASSIGN(auto arr, SparseUnionArray::create(*indices, {child1, child2}));
        ASSERT_OK(arr->validate_full());

        auto ty1 = STRUCT({field("a", int8()), field("b", uint16()), field("c", float32())});
        auto expected = assert_array_from_json(ty1, "[[0, 0, 0], [0, 65535, 1.5], [1, 42, -2.5]]");
        CheckView(arr, expected);
        CheckView(expected, arr);
    }

    TEST(TestArrayView, DecimalRoundTrip) {
        auto ty1 = decimal(10, 4);
        auto arr = assert_array_from_json(ty1, R"(["123.4567", "-78.9000", null])");

        auto ty2 = fixed_size_binary(16);
        ASSERT_OK_AND_ASSIGN(auto v, arr->View(ty2));
        ASSERT_OK(v->validate_full());
        ASSERT_OK_AND_ASSIGN(auto w, v->View(ty1));
        ASSERT_OK(w->validate_full());
        AssertArraysEqual(*arr, *w);
    }

    TEST(TestArrayView, Dictionaries) {
        // ARROW-6049
        auto ty1 = dictionary(int8(), float32());
        auto ty2 = dictionary(int8(), int32());

        auto indices = assert_array_from_json(int8(), "[0, 2, null, 1]");
        auto values = assert_array_from_json(float32(), "[0.0, 1.5, -2.5]");

        ASSERT_OK_AND_ASSIGN(auto expected_dict, values->View(int32()));
        ASSERT_OK_AND_ASSIGN(auto arr, DictionaryArray::from_arrays(ty1, indices, values));
        ASSERT_OK_AND_ASSIGN(auto expected,
                             DictionaryArray::from_arrays(ty2, indices, expected_dict));

        CheckView(arr, expected);
        CheckView(expected, arr);

        // Incompatible index type
        auto ty3 = dictionary(int16(), int32());
        CheckViewFails(arr, ty3);

        // Incompatible dictionary type
        auto ty4 = dictionary(int16(), float64());
        CheckViewFails(arr, ty4);

        // Check dictionary-encoded child
        auto offsets = assert_array_from_json(int32(), "[0, 2, 2, 4]");
        ASSERT_OK_AND_ASSIGN(auto list_arr, ListArray::from_arrays(*offsets, *arr));
        ASSERT_OK_AND_ASSIGN(auto expected_list_arr,
                             ListArray::from_arrays(*offsets, *expected));
        CheckView(list_arr, expected_list_arr);
        CheckView(expected_list_arr, list_arr);
    }

    TEST(TestArrayView, ExtensionType) {
        auto ty1 = std::make_shared<IPv4Type>();
        auto data = assert_array_from_json(ty1->storage_type(), R"(["ABCD", null])")->data();
        data->type = ty1;
        auto arr = ty1->make_array(data);
#if NEBULA_LITTLE_ENDIAN
        auto expected = assert_array_from_json(uint32(), "[1145258561, null]");
#else
        auto expected = assert_array_from_json(uint32(), "[1094861636, null]");
#endif
        CheckView(arr, expected);
        CheckView(expected, arr);
    }

    TEST(TestArrayView, NonZeroOffset) {
        auto arr = assert_array_from_json(int16(), "[10, 11, 12, 13]");

        ASSERT_OK_AND_ASSIGN(auto expected, arr->View(fixed_size_binary(2)));
        CheckView(arr->slice(1), expected->slice(1));
    }

    TEST(TestArrayView, NonZeroNestedOffset) {
        auto list_values = assert_array_from_json(int16(), "[10, 11, 12, 13, 14]");
        auto view_values = assert_array_from_json(uint16(), "[10, 11, 12, 13, 14]");

        auto list_offsets = assert_array_from_json(int32(), "[0, 2, 3]");

        ASSERT_OK_AND_ASSIGN(auto arr,
                             ListArray::from_arrays(*list_offsets, *list_values->slice(2)));
        ASSERT_OK_AND_ASSIGN(auto expected,
                             ListArray::from_arrays(*list_offsets, *view_values->slice(2)));
        ASSERT_OK(arr->validate_full());
        CheckView(arr->slice(1), expected->slice(1));

        // Be extra paranoid about checking offsets
        ASSERT_OK_AND_ASSIGN(auto result, arr->slice(1)->View(expected->type()));
        ASSERT_EQ(1, result->offset());
        ASSERT_EQ(2, static_cast<const ListArray &>(*result).values()->offset());
    }

}  // namespace nebula
